package wireguardcommon

import (
	"bytes"
	"crypto"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"fmt"
	"github.com/aead/ecdh"
	"github.com/vishvananda/netlink"
	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/curve25519"
	"golang.zx2c4.com/wireguard/wgctrl"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
	"io"
	"log"
	"net"
	"net/http"
	"net/netip"
	"os"
	"time"
)

const (
	InitialPort                 = 32323
	WireguardListenPort         = 32322
	PersistentKeepAliveInterval = 32
	InterfaceName               = "wg3"
	Ip2                         = "5.161.133.68"
	Ip1                         = "192.168.1.152"
	Ip3                         = "127.0.0.1"
	IPRange                     = "192.168.88.1/24"
)

const (
	ExitSetupSuccess = 0
	ExitSetupFailed  = 1
)

func GenerateKeyPair() (privateKey crypto.PrivateKey, publicKey crypto.PublicKey, err error) {
	filename := "resources/managerPrivateKey"
	if _, err = os.Stat(filename); errors.Is(err, os.ErrNotExist) {
		fmt.Printf("managerPrivateKey not found")
		c25519 := ecdh.X25519()
		privateKey, publicKey, err = c25519.GenerateKey(rand.Reader)
		if err != nil {
			fmt.Printf("Failed to generate manager privateKey/public key pair: %s\n", err)
			return
		}

		var managerPrivateBytes [32]byte
		if ok := CheckType(&managerPrivateBytes, privateKey); !ok {
			panic("ecdh: unexpected type of manager public key")
		}

		var managerPublicBytes [32]byte
		if ok := CheckType(&managerPublicBytes, publicKey); !ok {
			panic("ecdh: unexpected type of manager public key")
		}

		managerDHPrivate := base64.StdEncoding.EncodeToString(managerPrivateBytes[:])
		managerDHPublic := base64.StdEncoding.EncodeToString(managerPublicBytes[:])
		fmt.Printf("managerDHPrivate: %s\n", managerDHPrivate)
		fmt.Printf("managerDHPublic: %s\n", managerDHPublic)
		err := os.WriteFile(filename, []byte(managerDHPrivate), 0600)
		if err != nil {
			log.Fatalf("Write file for privateKey key failed, %s", err)
		}
	} else {
		fmt.Printf("managerPrivateKey found\n")
		var managerBytes []byte
		var managerPrivateBytes []byte
		var managerPublicBytes [32]byte
		managerBytes, err = os.ReadFile(filename)
		if err != nil {
			log.Fatalf("Read file for privateKey key failed, %s", err)
		}
		fmt.Printf("managerBytes = %s\n", managerBytes)
		managerPrivateBytes, err = base64.StdEncoding.DecodeString(string(managerBytes))
		privateKey, err = base64.StdEncoding.DecodeString(string(managerBytes))

		// build public key from privateKey key
		publicKey, err = curve25519.X25519(managerPrivateBytes, curve25519.Basepoint)
		if ok := CheckType(&managerPublicBytes, publicKey); !ok {
			panic("ecdh: unexpected type of manager public key")
		}

		managerDHPrivate := base64.StdEncoding.EncodeToString(managerPrivateBytes[:])
		managerDHPublic := base64.StdEncoding.EncodeToString(managerPublicBytes[:])
		fmt.Printf("managerDHPublic: %s\n", managerDHPublic)
		fmt.Printf("managerDHPrivate: %s\n", managerDHPrivate)
	}
	return
}

func CheckType(key *[32]byte, typeToCheck interface{}) (ok bool) {
	switch t := typeToCheck.(type) {
	case [32]byte:
		copy(key[:], t[:])
		ok = true
	case *[32]byte:
		copy(key[:], t[:])
		ok = true
	case []byte:
		if len(t) == 32 {
			copy(key[:], t)
			ok = true
		}
	case *[]byte:
		if len(*t) == 32 {
			copy(key[:], *t)
			ok = true
		}
	}
	return
}

func GenerateSignature(url string, message string) (signature string, err error) {
	if url == "" {
		url = "http://test1.jagat.me:8123/nacl-sign"
	}
	requestJson := "{\"message\": \"" + message + "\"}"
	response, err := http.Post(url, "application/json", bytes.NewBuffer([]byte(requestJson)))
	if err != nil {
		log.Fatalf("Cant reach the url for generating signature, %s", err)
	}
	signatureBytes, err := io.ReadAll(response.Body)
	signature = string(signatureBytes)
	return
}
func VerifySignature(url string, message string, signature string) (status bool, err error) {
	if url == "" {
		url = "http://test1.jagat.me:8123/nacl-verify"
	}
	status = false
	requestJson := "{\"message\": \"" + message + "\", \"signature\":\"" + signature + "\"}"
	response, err := http.Post(url, "application/json", bytes.NewBuffer([]byte(requestJson)))
	if err != nil {
		log.Fatalf("Cant reach the url for verifying signature, %s", err)
	}
	if response.StatusCode == 200 {
		status = true
		return
	}
	return
}

func GenPskOnEdge(managerDHPublicStr string) (presharedKey string, edgeDHPublic string, err error) {
	c25519 := ecdh.X25519()
	var edgePrivate crypto.PrivateKey
	var edgePublic crypto.PublicKey

	edgePrivate, edgePublic, err = c25519.GenerateKey(rand.Reader)

	var edgePublicBytes [32]byte
	if ok := CheckType(&edgePublicBytes, edgePublic); !ok {
		panic("ecdh: unexpected type of edge public key")
	}
	edgeDHPublic = base64.StdEncoding.EncodeToString(edgePublicBytes[:])
	//fmt.Printf("edgeDHPublic: %s\n", edgeDHPublic)

	if err != nil {
		fmt.Printf("Failed to generate edge private/public key pair: %s\n", err)
		return
	}

	managerDHPublic, err := base64.StdEncoding.DecodeString(managerDHPublicStr)
	if err != nil {
		fmt.Printf("manager public key is could not be decoded: %s\n", err)
	}

	if err = c25519.Check(managerDHPublic); err != nil {
		fmt.Printf("manager public key is not on the curve: %s\n", err)
	}
	secret := c25519.ComputeSecret(edgePrivate, managerDHPublic)
	presharedKey = base64.StdEncoding.EncodeToString(secret)
	return
}

func GenPskOnManager(edgeDHPublicStr string,
	managerPrivate crypto.PrivateKey) (presharedKey string, err error) {
	c25519 := ecdh.X25519()

	edgeDHPublic, err := base64.StdEncoding.DecodeString(edgeDHPublicStr)
	if err != nil {
		fmt.Printf("edge public key is could not be decoded: %s\n", err)
	}

	if err = c25519.Check(edgeDHPublic); err != nil {
		fmt.Printf("edge public key is not on the curve: %s\n", err)
	}
	secret := c25519.ComputeSecret(managerPrivate, edgeDHPublic)
	presharedKey = base64.StdEncoding.EncodeToString(secret)
	return
}

func WriteConn(conn net.Conn, message []byte) (writeLen int, err error) {
	// Send manager DH public
	writeLen, err = conn.Write(message)
	if err != nil {
		fmt.Printf("Errored when send data on %s, error=%d\n", conn, err)
		return
	}
	return
}
func ReadConn(conn net.Conn, maxLength int) (responseStr string, responseLen int, err error) {
	// Make a buffer to hold incoming data.
	edgeResponse := make([]byte, maxLength+32)
	// Read the incoming connection into the buffer.
	responseLen, err = conn.Read(edgeResponse)
	if err != nil {
		fmt.Println("Error reading:", err.Error())
		CloseConn(conn)
		return
	}
	if responseLen > maxLength {
		fmt.Printf("read more than needed. responseLen = %d\n", responseLen)
		CloseConn(conn)
		return
	}
	responseStr = string(edgeResponse[:responseLen])
	return
}

func CloseConn(conn net.Conn) {
	err := conn.Close()
	if err != nil {
		fmt.Printf("error closing conn. ignored\n")
		return
	}
}

func ResolveOutgoingInterface() string {
	routes, _ := netlink.RouteGet(net.ParseIP("1.1.1.1"))
	for _, route := range routes {
		log.Printf("connected on linkindexlink %d\n", route.LinkIndex)
		link, err := netlink.LinkByIndex(route.LinkIndex)
		if err != nil {
			fmt.Printf("unable to get the link. error = %s\n", err)
		} else {
			log.Printf("connected on olink %s\n", link.Attrs().Name)
			return link.Attrs().Name
		}
	}
	return ""
}

func ResolveHostIp() string {
	netInterfaceAddresses, err := net.InterfaceAddrs()

	if err != nil {
		return ""
	}

	for _, netInterfaceAddress := range netInterfaceAddresses {
		fmt.Printf("netInterfaceAddress = %s\n", netInterfaceAddress)
		networkIp, ok := netInterfaceAddress.(*net.IPNet)
		if ok && !networkIp.IP.IsLoopback() && networkIp.IP.To4() != nil {
			ip := networkIp.IP.String()
			fmt.Println("Resolved Host IP: " + ip)
			return ip
		}
	}
	return ""
}

// ExternalIP this does not list the wifi and bridge interfaces
func ExternalIP() (string, error) {
	var retIp net.IP
	ifaces, err := net.Interfaces()
	if err != nil {
		return "", err
	}
	for _, iface := range ifaces {
		fmt.Printf("iface=%s\n", iface)

		if iface.Flags&net.FlagUp == 0 {
			continue // interface down
		}
		if iface.Flags&net.FlagLoopback != 0 {
			continue // loopback interface
		}
		addrs, err := iface.Addrs()
		if err != nil {
			return "", err
		}
		for _, addr := range addrs {
			var ip net.IP
			switch v := addr.(type) {
			case *net.IPNet:
				ip = v.IP
			case *net.IPAddr:
				ip = v.IP
			}
			if ip == nil || ip.IsLoopback() {
				continue
			}
			ip = ip.To4()
			if ip == nil {
				continue // not an ipv4 address
			}
			fmt.Printf("ip = %s", ip.String())
			retIp = ip
			//return ip.String(), nil
		}
		return retIp.String(), nil
	}
	return "", errors.New("not connected to the network")
}

func CreateKeyPairFromSeed(seed string) (privateKeyStr string, publicKeyStr string, err error) {
	var edgePublicKey []byte
	seedHash := sha256.Sum256([]byte(seed))
	//create a curve25519 point (for privateKey)
	seedHash[0] &= 248
	seedHash[31] &= 127
	seedHash[31] |= 64
	privateKeyStr = base64.StdEncoding.EncodeToString(seedHash[:])
	fmt.Printf("privateKey = %s\n", privateKeyStr)

	// build publickey from privatekey
	edgePublicKey, err = curve25519.X25519(seedHash[:], curve25519.Basepoint)
	if err != nil {
		log.Fatalf("curve25519.X25519() failed: %v", err)
		return
	}
	publicKeyStr = base64.StdEncoding.EncodeToString(edgePublicKey)
	fmt.Printf("publicKey  = %s\n", publicKeyStr)

	return
}

func Chacha20poly1305Decrypt(key []byte, encryptedMsg []byte) (message []byte, err error) {
	var aead cipher.AEAD
	aead, err = chacha20poly1305.NewX(key)
	if err != nil {
		fmt.Printf("failed to create aead from key")
		return
	}
	if len(encryptedMsg) < aead.NonceSize() {
		panic("ciphertext too short")
	}

	// Split nonce and ciphertext.
	nonce, ciphertext := encryptedMsg[:aead.NonceSize()], encryptedMsg[aead.NonceSize():]

	// Decrypt the message and check it wasn't tampered with.
	message, err = aead.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return
	}
	return
}

func Chacha20poly1305Encrypt(key []byte, message []byte) (encryptedMsg []byte, err error) {
	var aead cipher.AEAD
	aead, err = chacha20poly1305.NewX(key)
	if err != nil {
		fmt.Printf("failed to create aead from key")
		return
	}
	// Select a random nonce, and leave capacity for the ciphertext.
	nonce := make([]byte, aead.NonceSize(), aead.NonceSize()+len(message)+aead.Overhead())
	if _, err = rand.Read(nonce); err != nil {
		return
	}

	// Encrypt the message and append the ciphertext to the nonce.
	encryptedMsg = aead.Seal(nonce, nonce, message, nil)
	return
}

type WireguardLink struct {
	netlink.LinkAttrs
}

func (generic *WireguardLink) Attrs() *netlink.LinkAttrs {
	return &generic.LinkAttrs
}

func (generic *WireguardLink) Type() string {
	return "wireguard"
}

func CreateWireguardInterface(privateKey crypto.PrivateKey, ipaddrAndMask string) (outgoingLink netlink.Link, wireguardLink netlink.Link, err error) {
	la := netlink.NewLinkAttrs()
	la.Name = InterfaceName
	linkExists := true
	link, err := netlink.LinkByName(la.Name)
	if err != nil || link == nil {
		linkExists = false
	}

	var outgoingInterfaceName string
	if !linkExists {
		fmt.Printf("interface %s does not exist. Creating\n", la.Name)
		link := WireguardLink{la}
		err = netlink.LinkAdd(&link)
		if err != nil {
			fmt.Printf("Failed to add link for %s, err: %s\n", la.Name, err)
			return nil, nil, err
		}
		outgoingInterfaceName = ResolveOutgoingInterface()
		wireguardLink, err = netlink.LinkByName(la.Name)
		// not needed to set master
		//netlink.LinkSetMaster(&link, outgoingLink)

		//////////////////////////////////////////////////////////////////////
		//configure wireguard interface with specific details
		//////////////////////////////////////////////////////////////////////
		err = AddLinkAddress(wireguardLink, ipaddrAndMask, "")
		if err != nil {
			fmt.Printf("Failed to set address to %s, err: %s\n", la.Name, err)
			return nil, nil, err
		}

		var privateBytes [32]byte
		if ok := CheckType(&privateBytes, privateKey); !ok {
			panic("ecdh: unexpected type of manager public key")
		}

		key, _ := wgtypes.NewKey(privateBytes[:])
		port := WireguardListenPort

		cfg := wgtypes.Config{
			PrivateKey:   &key,
			ListenPort:   &port,
			ReplacePeers: true,
		}

		var wgctrlClient *wgctrl.Client
		wgctrlClient, err = wgctrl.New()
		if err != nil {
			return
		}
		err = wgctrlClient.ConfigureDevice(la.Name, cfg)
		if err != nil {
			fmt.Println("wgctrlclient ConfigureDevice failed:", err.Error())
			return
		}
		//////////////////////////////////////////////////////////////////////
		//end of configure wireguard interface with manager specific details
		//////////////////////////////////////////////////////////////////////

	} else {
		fmt.Printf("interface %s exists\n", la.Name)
		outgoingInterfaceName = ResolveOutgoingInterface()
		wireguardLink, _ = netlink.LinkByName(la.Name)
	}

	// bring the interface up
	err = netlink.LinkSetUp(wireguardLink)
	if err != nil {
		fmt.Printf("Could not set the link up: %s\n", err)
		return nil, nil, err
	}
	outgoingLink, _ = netlink.LinkByName(outgoingInterfaceName)
	return
}

func AddPeer(presharedKey string, edgePublicStr string, endpointIp string, allowedIp string, replace bool) {
	ipNet, err := ConstructIPAndMask(allowedIp)
	if err != nil {
		fmt.Println("ParseCIDR() failed:", err.Error())
		return
	}
	var allowedIps []net.IPNet
	allowedIps = append(allowedIps, *ipNet)

	presharedKeyBytes, _ := base64.StdEncoding.DecodeString(presharedKey)
	edgePublicBytes, _ := base64.StdEncoding.DecodeString(edgePublicStr)

	key, _ := wgtypes.NewKey(edgePublicBytes)
	pk, _ := wgtypes.NewKey(presharedKeyBytes)
	var dur = time.Second * PersistentKeepAliveInterval
	peerConfig := wgtypes.PeerConfig{
		PublicKey:                   key,
		PresharedKey:                &pk,
		ReplaceAllowedIPs:           true,
		PersistentKeepaliveInterval: &dur,
		AllowedIPs:                  allowedIps,
	}

	var ipEndpoint net.IP
	if endpointIp != "" {
		ipEndpoint = net.ParseIP(endpointIp)
		peerConfig.Endpoint = &net.UDPAddr{
			IP:   ipEndpoint,
			Port: WireguardListenPort,
		}
	}

	var peers []wgtypes.PeerConfig
	peers = append(peers, peerConfig)
	cfg := wgtypes.Config{
		ReplacePeers: replace,
		Peers:        peers,
	}

	wgctrlClient, err := wgctrl.New()
	if err != nil {
		fmt.Println("wgctrl.New() failed:", err.Error())
		return
	}
	err = wgctrlClient.ConfigureDevice(InterfaceName, cfg)
	if err != nil {
		fmt.Println("wgctrlclient ConfigureDevice failed:", err.Error())
		return
	}
}

func AddLinkAddress(link netlink.Link, linkAddressAndMask string, peerAddress string) (err error) {
	//////////////////////////////////////////////////////////////////////
	//configure wireguard interface with manager specific details
	//////////////////////////////////////////////////////////////////////
	fmt.Printf("linkAddressAndMask=%s\n", linkAddressAndMask)
	var netmask *net.IPNet
	netmask, err = ConstructIPAndMask(linkAddressAndMask)
	var addr = &netlink.Addr{IPNet: netmask}

	if peerAddress != "" {
		fmt.Printf("peerAddress %s\n", peerAddress)
		peerAddr, _ := ConstructIPAndMask(peerAddress)
		fmt.Printf("peerAddr %s\n", peerAddr)
		addr.Peer = peerAddr
	}
	err = netlink.AddrAdd(link, addr)
	if err != nil {
		fmt.Printf("Failed to set address to %s, err: %s\n", InterfaceName, err)
		return
	}
	return
}

func ConstructIPAndMask(inputStr string) (netipnet *net.IPNet, err error) {
	fmt.Printf("inputStr=%s\n", inputStr)
	var netip net.IP
	netip, netipnet, err = net.ParseCIDR(inputStr)
	if err != nil {
		netip = net.ParseIP(inputStr)
		if netip == nil {
			fmt.Printf("Failed to parse address %s, err: %s\n", inputStr, err)
			return
		} else {
			netip, netipnet, _ = net.ParseCIDR(inputStr + "/32")
			err = nil
		}
	}
	netipnet.IP = netip
	return
}

func GetPeers() (deviceIpMap map[netip.Addr]wgtypes.Peer, err error) {
	deviceIpMap = make(map[netip.Addr]wgtypes.Peer)
	wgctrlClient, err := wgctrl.New()
	defaultDevice, err := wgctrlClient.Device(InterfaceName)
	if err != nil {
		fmt.Printf("Failed to find the default wireguard device %s, err: %s\n", InterfaceName, err)
		return
	}
	devices, err := wgctrlClient.Devices()
	if err != nil {
		return
	}
	for _, device := range devices {
		fmt.Printf("device = %s\n", device.Name)
		fmt.Printf("Peers:\n")
		for _, p := range device.Peers {
			fmt.Printf(" public=%s, psk=%s, ips=%s, endpoint=%s\n", p.PublicKey, p.PresharedKey, p.AllowedIPs, p.Endpoint)
		}
	}

	var peers = defaultDevice.Peers
	for _, p := range peers {
		allowedIp, err := netip.ParseAddr(p.AllowedIPs[0].IP.String())
		if err != nil {
			fmt.Printf("failed to parse %s\n", p.AllowedIPs[0].IP.String())
		}
		deviceIpMap[allowedIp] = p
	}

	return
}

func GetFirstAddr() (netip.Addr, *net.IPNet) {
	_, iPRange, _ := net.ParseCIDR(IPRange)
	ipFirst, _ := netip.ParseAddr(iPRange.IP.String())
	return ipFirst.Next(), iPRange
}

func FindUnusedIp() (ipNext netip.Addr, err error) {
	peers, _ := GetPeers()
	ipNext, ipRange := GetFirstAddr()

	ipNext = ipNext.Next()
	ip := net.ParseIP(ipNext.String())
	for {
		if _, ok := peers[ipNext]; !ok {
			if ipRange.Contains(ip) {
				return ipNext, nil
			}
		}
		ipNext = ipNext.Next()
		ip = net.ParseIP(ipNext.String())
	}
	return ipNext, errors.New("cant find unused ip address")
}
